/*
 * Copyright (c) China Telecom Cloud Technology Co., Ltd. 2024-2025. All rights reserved.
 */
#ifndef SEGMENT_H
#define SEGMENT_H

#include "KvcBaseDefine.h"
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>

#if defined(USE_GPU)
#include <cuda_runtime.h>
#elif defined(USE_NPU)
#include "acl/acl.h"
#endif

/**
 *! @brief   Segment, Segment对标LMCache的CHUNK
 *
 */
class Segment {
private:

    int32_t     _poolId;               // 归属Pool ID
    int64_t     _segmentId;            // 当前Segment ID
    uint64_t    _offset;               // 当前Segment相对于共享内存首地址的offset
    uint8_t*    _memAddr;              // 当前Segment的内存首地址
    int32_t     _memLength;            // 当前Segment的内存大小，单位Byte
    uint64_t    _accessTime;           // 当前Segment的最近访问时间
    bool        _isShare;              // 是否共享内存（预留）
    bool        _isFilled;             // 是否已填充数据(TP填满过后Segment就成为READ ONLY状态)
    std::string _shmFilePath;          // Segment所在的共享内存文件路径
    bool        _isSaved;              // 是否已持久化

    std::atomic<int> _writeInRecord;   // TP场景下Segment的写入次数(MLA除外)

    // CacheInfo 信息
    int32_t     _startLayer;
    int32_t     _endLayer;
    int32_t     _layerNum;
    int32_t     _layerCacheSize;

    // BlockHashId 信息
    int64_t     _hashId;                           // 当前segment对应的hashId，与_segmentId相同
    int64_t     _parentHashId;                     // 当前segment的前一个segment对应的hashId
    uint64_t    _localHashId;                      // 当前segment对应的位置无关hashId
    char        _verifyHashId[VERIFY_HASH_LENGTH]; // 当前segment的hash冲突时的备用hash

public:
    Segment() {}

    Segment(int64_t segment_id, 
        uint64_t    offset, 
        uint8_t *   memAddr,
        size_t      memLength, 
        bool        isFilled, 
        std::string shmFilePath,
        int32_t writeRecordSize = 0, 
        int32_t startLayer      = 0,
        int32_t endLayer        = 0, 
        int32_t layerNum        = 0,
        int32_t layerCacheSize  = 0, 
        int64_t parentHashId    = 0,
        uint64_t localHashId    = 0)
        : _poolId(0), 
        _segmentId(segment_id), 
        _offset(offset), 
        _memAddr(memAddr),
        _memLength(memLength), 
        _accessTime(0), 
        _isShare(true),
        _isFilled(isFilled), 
        _isSaved(false), 
        _shmFilePath(shmFilePath),
        _startLayer(startLayer), 
        _endLayer(endLayer), 
        _layerNum(layerNum),
        _layerCacheSize(layerCacheSize), 
        _hashId(segment_id),
        _parentHashId(parentHashId), 
        _localHashId(localHashId) {
        _writeInRecord.store(writeRecordSize);
    }


    Segment(const Segment &s) {
        init(s.getPoolId(), 
            s.getSegmentId(), 
            s.getOffset(), 
            s.getMemAddr(),
            s.getSegmentSize(), 
            s.getAccessTime(), 
            s.getIsShare(), 
            s.getIsFilled(),
            s.getShmFilePath(), 
            s.getIsSaved(), 
            s.getWriteInRecord(),
            s.getStartLayer(), 
            s.getEndLayer(), 
            s.getLayerNum(),
            s.getLayerCacheSize(), 
            s.getParentHashId(), 
            s.getLocalHashId());
        }

    void init(int32_t poolId, 
            int64_t segment_id, 
            uint64_t offset,
            uint8_t *memAddr, 
            size_t memLength, 
            uint64_t accessTime,
            bool isShare, 
            bool isFilled, 
            std::string shmFilePath, 
            bool isSaved,
            int32_t writeRecordSize = 0, 
            int32_t startLayer      = 0,
            int32_t endLayer        = 0, 
            int32_t layerNum        = 0,
            int32_t layerCacheSize  = 0, 
            int64_t parentHashId    = 0,
            uint64_t localHashId    = 0
        );

    /* 增加atomic成员后，需要重载= */
    void operator=(const Segment &s) {
        init(s.getPoolId(), 
            s.getSegmentId(), 
            s.getOffset(), 
            s.getMemAddr(),
            s.getSegmentSize(), 
            s.getAccessTime(), 
            s.getIsShare(), 
            s.getIsFilled(),
            s.getShmFilePath(), 
            s.getIsSaved(), 
            s.getWriteInRecord(),
            s.getStartLayer(), 
            s.getEndLayer(), 
            s.getLayerNum(),
            s.getLayerCacheSize()
        );
    }

    /**
     *! @brief 获取当前Segment归属Pool ID
     * 
     * @return int32_t 
     */
    int32_t getPoolId() const;


    /**
     *! @brief 获取当前Segment对象的Id
     * 
     * @return int64_t 
     */
    int64_t getSegmentId() const;


    /**
     *! @brief 获取当前Segment相对于共享内存首地址的offset
     * 
     * @return uint64_t 
     */
    uint64_t getOffset() const;


    /**
     *! @brief 获取当前Segment对象的数据首地址
     * 
     * @return uint8_t* 
     */
    uint8_t *getMemAddr() const;


    /**
     *! @brief 获取当前Segment的内存大小，单位Byte
     * 
     * @return int32_t 
     */
    int32_t getSegmentSize() const;


    /**
     *! @brief 获取当前Segment的最近访问时间
     * 
     * @return uint64_t 
     */
    uint64_t getAccessTime() const;


    /**
     *! @brief 当前Segment是否共享内存（预留）
     * 
     * @return bool 
     */
    bool getIsShare() const;


    /**
     *! @brief 当前Segment是否已经填充数据
     * 
     * @return bool
     */
    bool getIsFilled() const;


    /**
     *! @brief 获取Segment所在的共享内存文件路径
     * 
     * @return std::string 
     */
    std::string getShmFilePath() const;


    /**
     *! @brief 获取当前写入的数据次数(TP场景下一个Segment需要多次写入，MLA除外)
     * 
     * @return int32_t 
     */
    int32_t getWriteInRecord() const;


    /**
     *! @brief 获取当前Segment数据对应模型的起始层
     * 
     * @return int32_t 
     */
    int32_t getStartLayer() const;


    /**
     *! @brief 获取当前Segment数据对应模型的结束层
     * 
     * @return int32_t 
     */
    int32_t getEndLayer() const;


    /**
     *! @brief 获取当前Segment数据对应模型层数
     * 
     * @return int32_t 
     */
    int32_t getLayerNum() const;

    /**
     *! @brief 获取每层的KV Cache大小(字节)
     * 
     * @return int32_t 
     */
    int32_t getLayerCacheSize() const;


    /**
     *! @brief 获取当前Segment的前缀哈希
     * 
     * @return int64_t 
     */
    int64_t getHashId() const;


    /**
     *! @brief 获取当前Segment的前一个Segment哈希
     * 
     * @return int64_t 
     */
    int64_t getParentHashId() const;


    /**
     *! @brief 获取当前Segment的局部哈希(仅本块)
     * 
     * @return uint64_t 
     */
    uint64_t getLocalHashId() const;


    /**
     *! @brief 获取当前Segment的hash冲突时的备用hash
     * 
     * @return const char* 
     */
    const char *getVerifyHashId() const;



    /**
     *! @brief 当前Segment是否已经持久化
     * 
     * @return bool
     */
    bool getIsSaved() const;


    /**
     *! @brief 设置 Segment Id 
     * 
     * @param segmentId 
     */
    void setSegmentId(int64_t segmentId);


    /**
     *! @brief 设置当前Segemnt是否已经填充数据
     * 
     * @param val 
     */
    void setIsFilled(bool val);


    /**
     *! @brief 设置当前Segment的数据是否已经持久化
     * 
     * @param val 
     */
    void setIsSaved(bool val);


    /**
     *! @brief 设置当前Segment最近的访问时间
     * 
     * @param timestamp 
     */
    void setAccessTime(uint64_t timestamp);


    /**
     *! @brief 设置当前Segment写入的大小
     * 
     * @param val 
     */
    //
    void setWriteInRecord(int32_t val);


    /**
     *! @brief 写入记录值减1
     * 
     */
    void decWriteInRecord();


    /**
     *! @brief 设置KV Cache信息
     * 
     * @param cacheInfo 
     */
    void setCacheInfo(CacheInfo &cacheInfo);
    

    /**
     *! @brief 设置块的哈希信息
     * 
     * @param blockHashId 
     */
    void setBlockHashId(BlockHashId &blockHashId);


    /**
     *! @brief  将内存(CPU/GPU/NPU)的数据写入Segment
     * 
     * @param devType   设备类型，CPU/GPU/NPU
     * @param buffer    待写入数据在原设备上的地址
     * @param segOffset 相对于Segment首地址的offset
     * @param size      待写入的数据量(字节)
     * @param privData  私有数据，可选(stream的指针), 如果指定流, 则使用异步拷贝
     * 
     * @return int32_t  执行成功返回0，失败返回非0
     */
    int32_t writeIn(DeviceType devType, 
                    void *buffer, 
                    uint64_t segOffset,
                    uint32_t size, 
                    void *privData = nullptr
                );


    /**
     *! @brief  将Segment中的数据拷贝到内存(CPU/GPU/NPU)目标地址
     * 
     * @param devType    设备类型，CPU/GPU/NPU
     * @param buffer     目的设备上的缓冲区地址
     * @param segOffset  相对于segment首地址的offset
     * @param size       待写入的数据量(字节)
     * @param privData   私有数据，可选(stream的指针), 如果指定流, 则使用异步拷贝
     * 
     * @return int32_t  执行成功返回0，失败返回非0
     */
    int32_t copyTo(DeviceType devType, 
                void *buffer, 
                uint64_t segOffset,
                uint32_t size, 
                void *privData = nullptr
            );
};

#endif // SEGMENT_H